#coding=utf-8
import unreal


def CreatePoleCtrl(hierarchy_ctrl, hierarchy_info, component_info, head_key, middle_key, tail_key, scale_length):
    '''
    hierarchy_ctrl: hierarchy controller
    hierarchy_info: hierarchy
    component_info: component info
    head_key: first bone
    middle_key: second bone
    tail_key: third bone
    scale_length: distance
    return ctrl_key
    '''
    ctrl_key = CreateEulerTransformCtrl(hierarchy_ctrl, hierarchy_info, component_info)
    
    head_pos = unreal.Vector()
    middle_pos = unreal.Vector()
    tail_pos = unreal.Vector()
    head_transform = hierarchy_info.get_global_transform(head_key, True)
    middle_transform = hierarchy_info.get_global_transform(middle_key, True)
    tail_transform = hierarchy_info.get_global_transform(tail_key, True)
    
    head_pos = head_transform.get_editor_property("translation")
    middle_pos = middle_transform.get_editor_property("translation")
    tail_pos = tail_transform.get_editor_property("translation")
    
    h_m_vec = middle_pos.subtract(head_pos)
    h_t_vec = tail_pos.subtract(head_pos)
    
    h_t_length = h_t_vec.length()
    middle_vec1 = h_t_vec.normal().multiply_float(h_m_vec.dot(h_t_vec) / h_t_length)
    
    middle_vec2 = h_m_vec.subtract(middle_vec1).normal().multiply_float(scale_length)
    temp_vec = head_pos.add(middle_vec1)
    final_pos = temp_vec.add(middle_vec2)
    
    pole_transform = unreal.Transform(location=final_pos)
    hierarchy_info.set_control_offset_transform(ctrl_key, pole_transform, True)

    return ctrl_key


def CreateEulerTransformCtrl(hierarchy_ctrl, hierarchy_info, component_info):
    '''
    hierarchy_ctrl: hierarchy controller
    hierarchy_info: hierarchy
    component_info: component info
    return ctrl_key
    '''
    ctrl_name = component_info[0]
    ctrl_shape = component_info[1]
    ctrl_shape_offset = component_info[2]
    ctrl_shape_rotation = component_info[3]
    ctrl_size = component_info[4]
    ctrl_color = component_info[5]

    # set ctrl settings
    ctrl_settings = unreal.RigControlSettings()
    ctrl_settings.control_type = unreal.RigControlType.EULER_TRANSFORM
    ctrl_settings.shape_name = ctrl_shape
    ctrl_settings.shape_color = ctrl_color

    transform_value = unreal.EulerTransform(location=[0,0,0], rotation=[0,0,0], scale=[1,1,1])
    ctrl_value = hierarchy_info.make_control_value_from_euler_transform(transform_value)
    # create ctrl
    ctrl_key = hierarchy_ctrl.add_control(ctrl_name, "", ctrl_settings, ctrl_value)
    
    # set shape transform
    shape_transform = unreal.Transform(location=ctrl_shape_offset, rotation=ctrl_shape_rotation, scale=ctrl_size)
    hierarchy_info.set_control_shape_transform(ctrl_key, shape_transform, True)

    return ctrl_key


def CreateBoolCtrl(hierarchy_ctrl, hierarchy_info, ctrl_name, default_value):
    '''
    hierarchy_ctrl: hierarchy controller
    hierarchy_info: hierarchy
    ctrl_name: ctrl name
    default_value: default value
    return ctrl_key
    '''
    # set ctrl settings
    ctrl_settings = unreal.RigControlSettings()
    ctrl_settings.control_type = unreal.RigControlType.BOOL

    ctrl_value = hierarchy_info.make_control_value_from_bool(default_value)
    ctrl_key = hierarchy_ctrl.add_control(ctrl_name, "", ctrl_settings, ctrl_value)

    return ctrl_key


def CreateFloatCtrl(hierarchy_ctrl, hierarchy_info, ctrl_name, minimum_value, default_value, maximum_value, shape_visible, component_info=""):
    '''
    hierarchy_ctrl: hierarchy controller
    hierarchy_info: hierarchy
    ctrl_name: ctrl name
    minimum_value: minimum value
    default_value: default value
    maximum_value: maximum value
    shape_visible: shape visible
    component_info: component info, if shape_visible false, component_info set ""
    return ctrl_key
    '''
    # set ctrl settings
    ctrl_settings = unreal.RigControlSettings()
    ctrl_settings.control_type = unreal.RigControlType.FLOAT
    ctrl_settings.minimum_value = hierarchy_info.make_control_value_from_float(minimum_value)
    ctrl_settings.maximum_value = hierarchy_info.make_control_value_from_float(maximum_value)
    ctrl_settings.limit_enabled = [unreal.RigControlLimitEnabled(True, True)]

    ctrl_value = hierarchy_info.make_control_value_from_float(default_value)
    ctrl_key = unreal.RigElementKey()
    if shape_visible:
        ctrl_settings.shape_visible = shape_visible
        ctrl_settings.shape_name = component_info[1]
        ctrl_settings.shape_color = component_info[5]

        ctrl_shape_offset = component_info[2]
        ctrl_shape_rotation = component_info[3]
        ctrl_size = component_info[4]
        ctrl_name = component_info[0]

        ctrl_key = hierarchy_ctrl.add_control(ctrl_name, "", ctrl_settings, ctrl_value)

        # set shape transform
        shape_transform = unreal.Transform(location=ctrl_shape_offset, rotation=ctrl_shape_rotation, scale=ctrl_size)
        hierarchy_info.set_control_shape_transform(ctrl_key, shape_transform, True)

    else:
        ctrl_settings.shape_visible = shape_visible
        ctrl_key = hierarchy_ctrl.add_control(ctrl_name, "", ctrl_settings, ctrl_value)

    return ctrl_key


def CreateIntegerCtrl(hierarchy_ctrl, hierarchy_info, ctrl_name, minimum_value, default_value, maximum_value, shape_visible, component_info=""):
    '''
    hierarchy_ctrl: hierarchy controller
    hierarchy_info: hierarchy
    ctrl_name: ctrl name
    minimum_value: minimum value
    default_value: default value
    maximum_value: maximum value
    shape_visible: shape visible
    component_info: component info, if shape_visible false, component_info set ""
    return ctrl_key
    '''
    # set ctrl settings
    ctrl_settings = unreal.RigControlSettings()
    ctrl_settings.control_type = unreal.RigControlType.FLOAT
    ctrl_settings.minimum_value = hierarchy_info.make_control_value_from_float(minimum_value)
    ctrl_settings.maximum_value = hierarchy_info.make_control_value_from_float(maximum_value)
    ctrl_settings.limit_enabled = [unreal.RigControlLimitEnabled(True, True)]

    ctrl_value = hierarchy_info.make_control_value_from_float(default_value)
    ctrl_key = unreal.RigElementKey()
    if shape_visible:
        ctrl_settings.shape_visible = shape_visible
        ctrl_settings.shape_name = component_info[1]
        ctrl_settings.shape_color = component_info[5]

        ctrl_shape_offset = component_info[2]
        ctrl_shape_rotation = component_info[3]
        ctrl_size = component_info[4]
        ctrl_name = component_info[0]

        ctrl_key = hierarchy_ctrl.add_control(ctrl_name, "", ctrl_settings, ctrl_value)

        # set shape transform
        shape_transform = unreal.Transform(location=ctrl_shape_offset, rotation=ctrl_shape_rotation, scale=ctrl_size)
        hierarchy_info.set_control_shape_transform(ctrl_key, shape_transform, True)

    else:
        ctrl_settings.shape_visible = shape_visible
        ctrl_key = hierarchy_ctrl.add_control(ctrl_name, "", ctrl_settings, ctrl_value)

    return ctrl_key


def CreateGrp(hierarchy_ctrl, hierarchy_info, grp_name, target_key):
    '''
    hierarchy_ctrl: hierarchy controller
    hierarchy_info: hierarchy
    grp_name: group name
    target_key: set transform, if standard transform, target_key set ""
    return grp_key
    '''
    grp_transform = unreal.Transform()
    if target_key == "":
        grp_transform = unreal.Transform(location=[0,0,0], rotation=[0,0,0], scale=[1,1,1])
    else:
        grp_transform = hierarchy_info.get_global_transform(target_key, True)

    grp_key = hierarchy_ctrl.add_null(grp_name, "", grp_transform, True)

    return grp_key


def Set_Relative_Offset_Transform(hierarchy_info, target_key, parent_key):
    '''
    hierarchy_info: hierarchy
    target_key: the transform obj
    parent_key: the parent hierarchy
    '''
    target_transform = hierarchy_info.get_global_transform(target_key, True)
    parent_transform = hierarchy_info.get_global_transform(parent_key, True)

    relative_transform = target_transform.make_relative(parent_transform)

    hierarchy_info.set_control_offset_transform(target_key, relative_transform, True)


def Set_Offset_Transform(hierarchy_info, ctrl_key, bone_key, component_info):
    orient_info = component_info[6]
    bone_transform = hierarchy_info.get_global_transform(bone_key, True)
    if orient_info:
        bone_transform.set_editor_property("rotation", unreal.Quat(0,0,0,1))
    
    hierarchy_info.set_control_offset_transform(ctrl_key, bone_transform, True)


def CreateNode(graph_ctrl, node_lib, node_type, node_name):
    '''
    graph_ctrl: RigVMController
    node_lib: node_lib
    node_name: node_name
    return RigVMUnitNode
    '''
    script_struct_path = node_lib[node_type]
    cus_node_name = node_name
    rig_node = graph_ctrl.add_unit_node_from_struct_path(script_struct_path, node_name=cus_node_name)

    return rig_node

